from type import WindBreakInfo, WindBreakInfoDetail
from prompt import prompt, message
from langchain_community.chat_models import ChatZhipuAI
import re
import json
def extract_json(raw_text):
    # 匹配三个反引号包裹的JSON（常见于Markdown响应）
    match = re.search(r'```(?:json)?\n(.*?)\n```', raw_text, re.DOTALL)
    if match:
        raw_json = match.group(1)
    else:
        # 如果没有代码块标记，尝试直接提取最外层的大括号内容
        match = re.search(r'\{.*\}', raw_text, re.DOTALL)
        if not match:
            raise ValueError("No JSON found in the response")
        raw_json = match.group(0)

    try:
        return json.loads(raw_json)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON: {e}\nExtracted content: {raw_json}")
llm = ChatZhipuAI(
    api_key="49ec86958bf242e7b602e3636a72e812.mgDAyyf1gfaIzL9M",
    model="glm-4-plus",
    temperature=0.1,
)
model = llm.with_structured_output(WindBreakInfoDetail)
message = message.format(year = "2013",location = "武汉",ecosystem_name = "湿地生态系统")
lm = lm = [{"name": "planner", "role": "system", "content": prompt},
      {"name": "planner", "role": "human", "content": message}]
response =  model.invoke(lm)
print(response)
